#include "DictionaryLearner.hpp"

namespace SparseCoding
{
	const double DictionaryLearner::DEFAULT_STEP = 1.0;
	const double DictionaryLearner::DEFAULT_ISTA_STEP = 0.1;
	const double DictionaryLearner::DEFAULT_NOISE_VARIANCE = 0.1;
	const double DictionaryLearner::DEFAULT_TOLERANCE = 0.1;
	const std::string DictionaryLearner::DICTIONARY_OUTPUT_FILE = "dictionary.txt";
	const std::string DictionaryLearner::DICTIONARY_NAME = "dictionary";
	const std::string DictionaryLearner::OUTPUT_IMAGE_FILE = "basisfunctions.png";
	const std::string DictionaryLearner::INSTRUCTIONS = "Press escape at any time to quit and save the dictionary.";
	const std::string DictionaryLearner::TRAINING_COMPLETE = "Training complete.";

	DictionaryLearner::DictionaryLearner(std::vector<cv::Mat>& input, double basisFunctionCount, unsigned int patchLength, unsigned int batchCount)
	{
		visualizationInterpolation = cv::INTER_AREA;
		visualizationSize = cv::Size(0, 0);

		images = input;
		patchSize = patchLength;
		sampleCount = batchCount;

		step = DEFAULT_STEP;
		istaStep = DEFAULT_ISTA_STEP;
		noiseVariance = DEFAULT_NOISE_VARIANCE;
		tolerance = DEFAULT_TOLERANCE;

		//Intialize the dictionary to random values between -0.5 and 0.5.
		std::mt19937 randGen(time(nullptr));
		std::uniform_real_distribution<double> dist(-0.5, 0.5);

		dictionary = cv::Mat(patchSize * patchSize, basisFunctionCount, CV_64FC1);

		for (unsigned int i = 0; i < dictionary.rows; i++)
		{
			for (unsigned int j = 0; j < dictionary.cols; j++)
				dictionary.at<double>(i, j) = dist(randGen);
		}

		//Make the basis functions of unit norm.
		for (unsigned int i = 0; i < basisFunctionCount; i++)
			dictionary.col(i) /= cv::norm(dictionary.col(i));
	}

	//Add the L1 penalty to the sparse codes or clamp them to zero.
	void DictionaryLearner::Shrink(cv::Mat sparseCodes, double step, double noiseVariance)
	{
		double shrinkStep = step * noiseVariance;

		for (unsigned int i = 0; i < sparseCodes.rows; i++)
		{
			int codeSign = Sign(sparseCodes.at<double>(i, 0));

			//Result of applying the L1 penalty.
			double result = sparseCodes.at<double>(i, 0) - (double)(shrinkStep * codeSign);

			//Clamp to 0 if the sign of the code changed.
			if (std::abs(result) < noiseVariance)
				sparseCodes.at<double>(i, 0) = 0.0;
			else
				sparseCodes.at<double>(i, 0) = result;
		}
	}

	int DictionaryLearner::Sign(double d)
	{
		if (d >= 0.0)
			return 1;
		else
			return -1;
	}


	cv::Mat DictionaryLearner::ISTA(cv::Mat samples, cv::Mat basisFunctions, double step, double noiseVariance, double tolerance)
	{
		//Initialize sparse codes to zero.
		cv::Mat sparseCodes = cv::Mat::zeros(basisFunctions.cols, samples.cols, CV_64FC1);

		for (unsigned int i = 0; i < samples.cols; i++)
		{
			while (true)
			{
				//Magnitude used for checking for convergence.
				double originalMagnitude = cv::norm(sparseCodes.col(i));

				//Compute the reconstruction error of the input.
				cv::Mat reconstError = samples.col(i) - basisFunctions * sparseCodes.col(i);

				sparseCodes.col(i) += step * basisFunctions.t() * reconstError;

				//Add the L1 penalty or clamp to zero if necessary.
				Shrink(sparseCodes.col(i), step, noiseVariance);

				double newMagnitude = cv::norm(sparseCodes.col(i));

				double magnitudeChange = newMagnitude - originalMagnitude;

				if (std::abs(magnitudeChange) < tolerance)
					break;
			}
		}

		return sparseCodes;
	}

	cv::Mat DictionaryLearner::Visualize(cv::Mat& basisFunctions, unsigned int border, unsigned int patchSize)
	{
		int imagesPerRow = 0;
		int imagesPerCol = 0;

		if (std::pow(std::floor(std::sqrt(basisFunctions.cols)), 2.0) != basisFunctions.cols)
		{
			imagesPerRow = std::sqrt(basisFunctions.cols);
			imagesPerCol = static_cast<int>(std::ceil(static_cast<float>(basisFunctions.cols) / static_cast<float>(imagesPerRow)));
		}
		else
			imagesPerRow = imagesPerCol = std::sqrt(basisFunctions.cols);

		int imageRows = border + imagesPerRow * (border + patchSize);
		int imageCols = border + imagesPerCol * (border + patchSize);

		cv::Mat image = cv::Mat::zeros(imageRows, imageCols, CV_64FC1);

		unsigned int k = 0;

		for (unsigned int i = 0; i < imagesPerRow; i++)
		{
			for (unsigned int j = 0; j < imagesPerCol; j++)
			{
				cv::Mat reshaped = basisFunctions.col(k).clone().reshape(1, patchSize);

				double minValue = 0.0;

				for (unsigned int a = 0; a < reshaped.rows; a++)
				{
					for (unsigned int b = 0; b < reshaped.cols; b++)
					{
						double currentVal = reshaped.at<double>(a, b);

						if (minValue > currentVal)
							minValue = currentVal;
					}
				}

				reshaped -= minValue;

				double maxValue = 0.0;

				for (unsigned int a = 0; a < reshaped.rows; a++)
				{
					for (unsigned int b = 0; b < reshaped.cols; b++)
					{
						double currentVal = reshaped.at<double>(a, b);

						if (maxValue < currentVal)
							maxValue = currentVal;
					}
				}

				reshaped /= maxValue;

				unsigned int row = border + i * (patchSize + border);
				unsigned int col = border + j * (patchSize + border);
				reshaped.copyTo(image(cv::Rect(col, row, patchSize, patchSize)));

				k++;

				if (k == basisFunctions.cols)
					break;
			}
		}

		return image;
	}

	void DictionaryLearner::SetParameters(double _step, double _istaStep, double _noiseVariance)
	{
		step = _step;
		istaStep = _istaStep;
		noiseVariance = _noiseVariance;
	}

	void DictionaryLearner::SaveDictionary()
	{
		cv::FileStorage file(DICTIONARY_OUTPUT_FILE, cv::FileStorage::WRITE);
		file << DICTIONARY_NAME << dictionary;
	}

	void DictionaryLearner::SaveVisualization(cv::Size newSize, int interpolationType)
	{
		cv::Mat visualization = Visualize(dictionary, BORDER, patchSize);

		cv::resize(visualization, visualization, newSize, 0.0, 0.0, interpolationType);
		visualization.convertTo(visualization, CV_8UC1, 255.0);
		cv::imwrite(OUTPUT_IMAGE_FILE, visualization);
	}

	void DictionaryLearner::Train(unsigned int iterations, bool visualize)
	{
		//Random distribution for the images.
		std::mt19937 randGen(time(nullptr));
		std::uniform_int_distribution<unsigned int> imageDist(0, images.size() - 1);

		cv::Mat visualization;

		std::cout << INSTRUCTIONS << std::endl;

		for (unsigned int i = 0; i < iterations; i++)
		{
			unsigned int imageIndex = imageDist(randGen);
			cv::Mat currentImage = images[imageIndex];

			//Choose random patches from a randomly selected image.
			cv::Mat patches = ChooseRandomPatches(currentImage);

			//Infer the sparse coefficients for every image patch.
			cv::Mat coefficients = ISTA(patches, dictionary, istaStep, noiseVariance, tolerance);

			//Compute the reconstruction error for the dictionary update rule.
			cv::Mat reconstError = patches - dictionary * coefficients;

			//The dictionary is updated with a scalar multiple of it's average gradient.
			cv::Mat averageDictionaryGradient = cv::Mat::zeros(dictionary.size(), dictionary.type());

			//Compute the average gradient.
			for (unsigned int j = 0; j < patches.cols; j++)
				averageDictionaryGradient += reconstError.col(j) * coefficients.col(j).t();

			averageDictionaryGradient /= static_cast<double>(patches.cols);

			//Take one step in the direction of the average gradient.
			dictionary += step * averageDictionaryGradient;

			//Normalize the basis functions to keep them from growing unbounded.
			for (unsigned int j = 0; j < dictionary.cols; j++)
				dictionary.col(j) /= cv::norm(dictionary.col(j));

			if (visualize)
			{
				int waitTime = DISPLAY_DELAY;

				//How long for the opencv imshow window to wait.
				//Small values seem to not let some of the first images display.
				if (i == (iterations - 1))
				{
					waitTime = 0;
					std::cout << TRAINING_COMPLETE << std::endl;
				}

				visualization = Visualize(dictionary, BORDER, patchSize);

				if (visualizationSize.width != 0)
					cv::resize(visualization, visualization, visualizationSize, 0.0, 0.0, visualizationInterpolation);

				cv::imshow("basis functions", visualization);

				char key = cv::waitKey(waitTime);

				if (key == ESCAPE_KEY)
					break;
			}
		}
	}

	cv::Mat DictionaryLearner::ChooseRandomPatches(cv::Mat image)
	{
		std::mt19937 randGen(time(nullptr));
		std::uniform_real_distribution<double> dist(0.0, 1.0);

		cv::Mat patches = cv::Mat(dictionary.rows, sampleCount, CV_64FC1);

		unsigned int lastRow = image.rows - patchSize;
		unsigned int lastCol = image.cols - patchSize;

		for (unsigned int j = 0; j < sampleCount; j++)
		{
			unsigned int row = static_cast<unsigned int>(ceil(dist(randGen) * lastRow));
			unsigned int col = static_cast<unsigned int>(ceil(dist(randGen) * lastCol));

			cv::Mat subMatrix = image(cv::Rect(row, col, patchSize, patchSize));
			subMatrix = subMatrix.clone().reshape(1, dictionary.rows);

			subMatrix.copyTo(patches.col(j));
		}

		return patches;
	}
}
